"""
预处理Kvasir-SEG数据集
消化道息肉分割数据集处理
"""

import os
from glob import glob
import cv2
import numpy as np
from tqdm import tqdm
from sklearn.model_selection import train_test_split


def preprocess_kvasir(source_dir, output_size=256, test_split=0.2, use_full_dataset=True):
    """
    预处理Kvasir-SEG数据集
    
    Args:
        source_dir: 'Kvasir-SEG' 或 'sessile-main-Kvasir-SEG'
        output_size: 输出图像尺寸
        test_split: 验证集比例
        use_full_dataset: 是否使用完整数据集（True）或sessile子集（False）
    """
    
    # 输入路径
    if use_full_dataset:
        input_dir = f'inputs/{source_dir}'
        output_name = f'kvasir_{output_size}'
    else:
        input_dir = f'inputs/{source_dir}'
        output_name = f'kvasir_sessile_{output_size}'
    
    img_dir = os.path.join(input_dir, 'images')
    mask_dir = os.path.join(input_dir, 'masks')
    
    # 输出路径
    output_dir = f'inputs/{output_name}'
    os.makedirs(f'{output_dir}/images', exist_ok=True)
    os.makedirs(f'{output_dir}/masks/0', exist_ok=True)
    
    print("="*80)
    print(f"预处理 {source_dir} 数据集")
    print(f"目标尺寸: {output_size}×{output_size}")
    print("="*80)
    
    # 获取所有图像
    image_paths = sorted(glob(os.path.join(img_dir, '*.jpg')))
    
    if len(image_paths) == 0:
        print(f"错误: 在 {img_dir} 中找不到图像!")
        return
    
    print(f"\n找到 {len(image_paths)} 张图像")
    
    # 处理所有图像
    processed_count = 0
    skipped_count = 0
    
    for img_path in tqdm(image_paths, desc="处理中"):
        # 获取文件名（不含扩展名）
        basename = os.path.basename(img_path).replace('.jpg', '')
        
        # 对应的mask文件
        mask_path = os.path.join(mask_dir, f'{basename}.jpg')
        
        # 检查mask是否存在
        if not os.path.exists(mask_path):
            skipped_count += 1
            continue
        
        try:
            # 读取图像和mask
            img = cv2.imread(img_path)
            mask = cv2.imread(mask_path, cv2.IMREAD_GRAYSCALE)
            
            if img is None or mask is None:
                skipped_count += 1
                continue
            
            # 确保图像是RGB格式
            if len(img.shape) == 2:
                img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
            elif img.shape[2] == 4:
                img = cv2.cvtColor(img, cv2.COLOR_BGRA2BGR)
            
            # 调整大小
            img_resized = cv2.resize(img, (output_size, output_size), 
                                    interpolation=cv2.INTER_LINEAR)
            mask_resized = cv2.resize(mask, (output_size, output_size), 
                                     interpolation=cv2.INTER_NEAREST)
            
            # 二值化mask（确保只有0和255）
            mask_binary = (mask_resized > 127).astype(np.uint8) * 255
            
            # 保存
            cv2.imwrite(f'{output_dir}/images/{basename}.png', img_resized)
            cv2.imwrite(f'{output_dir}/masks/0/{basename}.png', mask_binary)
            
            processed_count += 1
            
        except Exception as e:
            print(f"\n处理 {basename} 时出错: {e}")
            skipped_count += 1
    
    print(f"\n处理完成:")
    print(f"  - 成功处理: {processed_count} 张")
    print(f"  - 跳过: {skipped_count} 张")
    print(f"  - 输出目录: {output_dir}")
    
    # 数据集统计
    print(f"\n数据集统计:")
    print(f"  - 总样本数: {processed_count}")
    print(f"  - 建议训练集: {int(processed_count * (1 - test_split))} 张")
    print(f"  - 建议验证集: {int(processed_count * test_split)} 张")
    
    # 生成训练命令
    print(f"\n" + "="*80)
    print("训练命令:")
    print("="*80)
    
    print(f"\n# EnhancedNestedUNet (主实验)")
    print(f"python train.py \\")
    print(f"    --dataset {output_name} \\")
    print(f"    --arch EnhancedNestedUNet \\")
    print(f"    --name Kvasir_EnhancedUNet \\")
    print(f"    --img_ext .png --mask_ext .png \\")
    print(f"    --input_w {output_size} --input_h {output_size} \\")
    print(f"    --epochs 150 --batch_size 8 --lr 0.0001 \\")
    print(f"    --deep_supervision True")
    
    print(f"\n# NestedUNet Baseline (对比实验)")
    print(f"python train.py \\")
    print(f"    --dataset {output_name} \\")
    print(f"    --arch NestedUNet \\")
    print(f"    --name Kvasir_NestedUNet \\")
    print(f"    --img_ext .png --mask_ext .png \\")
    print(f"    --input_w {output_size} --input_h {output_size} \\")
    print(f"    --epochs 150 --batch_size 8 --lr 0.0001")
    
    print(f"\n# U-Net Baseline (对比实验)")
    print(f"python train.py \\")
    print(f"    --dataset {output_name} \\")
    print(f"    --arch UNet \\")
    print(f"    --name Kvasir_UNet \\")
    print(f"    --img_ext .png --mask_ext .png \\")
    print(f"    --input_w {output_size} --input_h {output_size} \\")
    print(f"    --epochs 150 --batch_size 8 --lr 0.0001")
    
    return processed_count


def main():
    import argparse
    
    parser = argparse.ArgumentParser(description='预处理Kvasir-SEG数据集')
    parser.add_argument('--dataset', type=str, default='Kvasir-SEG',
                       choices=['Kvasir-SEG', 'sessile-main-Kvasir-SEG'],
                       help='数据集名称')
    parser.add_argument('--size', type=int, default=256,
                       help='输出图像尺寸 (默认: 256)')
    parser.add_argument('--test_split', type=float, default=0.2,
                       help='验证集比例 (默认: 0.2)')
    
    args = parser.parse_args()
    
    use_full = (args.dataset == 'Kvasir-SEG')
    
    preprocess_kvasir(
        source_dir=args.dataset,
        output_size=args.size,
        test_split=args.test_split,
        use_full_dataset=use_full
    )


if __name__ == '__main__':
    main()






